#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
血管分割模型架构可视化脚本
基于UNet的编码器-解码器模型
"""

import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.patches import FancyBboxPatch, Rectangle
import numpy as np

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

def draw_model_architecture():
    """绘制UNet血管分割模型架构图"""
    
    fig, ax = plt.subplots(1, 1, figsize=(16, 12))
    ax.set_xlim(0, 16)
    ax.set_ylim(0, 12)
    ax.axis('off')
    
    # 颜色定义
    colors = {
        'input': '#E8F4FD',      # 浅蓝色 - 输入
        'conv': '#FFE4B5',       # 浅橙色 - 卷积层
        'pool': '#FFB6C1',       # 浅粉色 - 池化层
        'upsample': '#98FB98',   # 浅绿色 - 上采样
        'concat': '#DDA0DD',     # 浅紫色 - 拼接
        'output': '#F0E68C'      # 浅黄色 - 输出
    }
    
    # 绘制标题
    ax.text(8, 11.5, 'UNet血管分割模型架构', fontsize=20, fontweight='bold', ha='center')
    ax.text(8, 11, 'EncoderDecoder + UNet Backbone + MultiOutputHead', fontsize=14, ha='center', style='italic')
    
    # 输入层
    input_box = FancyBboxPatch((1, 9.5), 2, 1, boxstyle="round,pad=0.1", 
                              facecolor=colors['input'], edgecolor='black', linewidth=2)
    ax.add_patch(input_box)
    ax.text(2, 10, '输入图像\n512×512×3', ha='center', va='center', fontsize=10, fontweight='bold')
    
    # 数据预处理
    preprocess_box = FancyBboxPatch((4, 9.5), 2.5, 1, boxstyle="round,pad=0.1", 
                                   facecolor='#F5F5DC', edgecolor='black', linewidth=1)
    ax.add_patch(preprocess_box)
    ax.text(5.25, 10, '数据预处理\nNormalization\nResize', ha='center', va='center', fontsize=9)
    
    # 编码器路径 (下采样)
    encoder_stages = [
        {'pos': (1, 8), 'size': (2, 0.8), 'text': 'Conv Block 1\n64 channels\n2 convs'},
        {'pos': (1, 6.5), 'size': (2, 0.8), 'text': 'Conv Block 2\n128 channels\n2 convs'},
        {'pos': (1, 5), 'size': (2, 0.8), 'text': 'Conv Block 3\n256 channels\n2 convs'},
        {'pos': (1, 3.5), 'size': (2, 0.8), 'text': 'Conv Block 4\n512 channels\n2 convs'},
        {'pos': (1, 2), 'size': (2, 0.8), 'text': 'Bottleneck\n1024 channels\n2 convs'}
    ]
    
    # 绘制编码器
    for i, stage in enumerate(encoder_stages):
        # 卷积块
        conv_box = FancyBboxPatch(stage['pos'], stage['size'][0], stage['size'][1], 
                                 boxstyle="round,pad=0.05", facecolor=colors['conv'], 
                                 edgecolor='black', linewidth=1)
        ax.add_patch(conv_box)
        ax.text(stage['pos'][0] + stage['size'][0]/2, stage['pos'][1] + stage['size'][1]/2, 
               stage['text'], ha='center', va='center', fontsize=8)
        
        # 池化层 (除了最后一层)
        if i < len(encoder_stages) - 1:
            pool_box = Rectangle((3.5, stage['pos'][1] - 0.5), 1, 0.5, 
                               facecolor=colors['pool'], edgecolor='black', linewidth=1)
            ax.add_patch(pool_box)
            ax.text(4, stage['pos'][1] - 0.25, 'MaxPool\n2×2', ha='center', va='center', fontsize=7)
            
            # 下采样箭头
            ax.arrow(4, stage['pos'][1] - 0.5, 0, -0.7, head_width=0.1, head_length=0.1, 
                    fc='red', ec='red')
    
    # 解码器路径 (上采样)
    decoder_stages = [
        {'pos': (13, 3.5), 'size': (2, 0.8), 'text': 'Conv Block 4\n512 channels\n2 convs'},
        {'pos': (13, 5), 'size': (2, 0.8), 'text': 'Conv Block 3\n256 channels\n2 convs'},
        {'pos': (13, 6.5), 'size': (2, 0.8), 'text': 'Conv Block 2\n128 channels\n2 convs'},
        {'pos': (13, 8), 'size': (2, 0.8), 'text': 'Conv Block 1\n64 channels\n2 convs'}
    ]
    
    # 绘制解码器
    for i, stage in enumerate(decoder_stages):
        # 上采样
        upsample_box = Rectangle((11.5, stage['pos'][1]), 1, 0.8, 
                               facecolor=colors['upsample'], edgecolor='black', linewidth=1)
        ax.add_patch(upsample_box)
        ax.text(12, stage['pos'][1] + 0.4, 'UpSample\n2×2', ha='center', va='center', fontsize=7)
        
        # 卷积块
        conv_box = FancyBboxPatch(stage['pos'], stage['size'][0], stage['size'][1], 
                                 boxstyle="round,pad=0.05", facecolor=colors['conv'], 
                                 edgecolor='black', linewidth=1)
        ax.add_patch(conv_box)
        ax.text(stage['pos'][0] + stage['size'][0]/2, stage['pos'][1] + stage['size'][1]/2, 
               stage['text'], ha='center', va='center', fontsize=8)
        
        # 跳跃连接
        skip_y = encoder_stages[len(encoder_stages) - 2 - i]['pos'][1] + 0.4
        ax.arrow(3, skip_y, 8.5, stage['pos'][1] + 0.4 - skip_y, 
                head_width=0.1, head_length=0.2, fc='blue', ec='blue', alpha=0.7,
                linestyle='--', linewidth=2)
        
        # 拼接操作
        concat_box = Rectangle((10, stage['pos'][1] + 0.2), 1, 0.4, 
                             facecolor=colors['concat'], edgecolor='black', linewidth=1)
        ax.add_patch(concat_box)
        ax.text(10.5, stage['pos'][1] + 0.4, 'Concat', ha='center', va='center', fontsize=7)
    
    # 解码头 (MultiOutputHead)
    decode_head_box = FancyBboxPatch((8, 1), 3, 1.2, boxstyle="round,pad=0.1", 
                                    facecolor='#FFDAB9', edgecolor='black', linewidth=2)
    ax.add_patch(decode_head_box)
    ax.text(9.5, 1.6, 'MultiOutputHead\n64→2 classes\nDice Loss', ha='center', va='center', 
           fontsize=10, fontweight='bold')
    
    # 输出层
    output_box = FancyBboxPatch((13, 0.2), 2, 0.8, boxstyle="round,pad=0.1", 
                               facecolor=colors['output'], edgecolor='black', linewidth=2)
    ax.add_patch(output_box)
    ax.text(14, 0.6, '输出分割图\n512×512×2\n(背景/血管)', ha='center', va='center', 
           fontsize=10, fontweight='bold')
    
    # 绘制主要的数据流箭头
    # 输入到编码器
    ax.arrow(2, 9.5, 0, -1, head_width=0.1, head_length=0.1, fc='green', ec='green', linewidth=2)
    
    # 瓶颈到解码器
    ax.arrow(3, 2.4, 8.5, 1.5, head_width=0.1, head_length=0.2, fc='green', ec='green', linewidth=2)
    
    # 解码器到输出头
    ax.arrow(14, 8, 0, -5.5, head_width=0.1, head_length=0.2, fc='green', ec='green', linewidth=2)
    ax.arrow(9.5, 2.2, 0, -1, head_width=0.1, head_length=0.1, fc='green', ec='green', linewidth=2)
    
    # 输出头到最终输出
    ax.arrow(11, 1.6, 2, -1, head_width=0.1, head_length=0.1, fc='green', ec='green', linewidth=2)
    
    # 添加图例
    legend_elements = [
        patches.Patch(color=colors['input'], label='输入层'),
        patches.Patch(color=colors['conv'], label='卷积块'),
        patches.Patch(color=colors['pool'], label='池化层'),
        patches.Patch(color=colors['upsample'], label='上采样'),
        patches.Patch(color=colors['concat'], label='特征拼接'),
        patches.Patch(color=colors['output'], label='输出层')
    ]
    ax.legend(handles=legend_elements, loc='upper right', bbox_to_anchor=(0.98, 0.98))
    
    # 添加模型参数信息
    model_info = """
模型配置参数:
• 输入尺寸: 512×512×3
• 基础通道数: 64
• 编码器阶段: 5层
• 解码器阶段: 4层
• 跳跃连接: 4个
• 损失函数: Dice Loss
• 类别数: 2 (背景/血管)
• 标准化: SyncBN
• 激活函数: ReLU
    """
    ax.text(0.5, 5.5, model_info, fontsize=9, verticalalignment='top', 
           bbox=dict(boxstyle="round,pad=0.5", facecolor='lightgray', alpha=0.8))
    
    plt.tight_layout()
    return fig

def save_architecture_diagram():
    """保存模型架构图"""
    fig = draw_model_architecture()
    
    # 保存为高分辨率图片
    plt.savefig('vessel_unet_architecture.png', dpi=300, bbox_inches='tight', 
                facecolor='white', edgecolor='none')
    plt.savefig('vessel_unet_architecture.pdf', bbox_inches='tight', 
                facecolor='white', edgecolor='none')
    
    print("模型架构图已保存:")
    print("- vessel_unet_architecture.png (PNG格式)")
    print("- vessel_unet_architecture.pdf (PDF格式)")
    
    plt.show()

if __name__ == "__main__":
    print("正在生成血管分割UNet模型架构图...")
    save_architecture_diagram()